import glob

from fun_wav import *


'''
这个文件主要是把所有信号处理的函数整合在一起显示，显示流程如下
时频域操作
1. 原始信号加滤波
2. FFT信号
3. 时频图和信号梯度图
4. 频率变化图
5. 曲线归一化图

微多普勒操作
1。 原始信号加滤波
2. 调制解调后的唇动信号
'''

def ShowSinFFT(Zxx, f, t):
    timeLen = Zxx.shape[1]
    for i in range(timeLen):
        freq = Zxx[:, i]
        plt.title(t[i])
        plt.plot(f, freq)
        plt.show()

def LowFiltWav(signalWav, N, lowPass):
    b, a = signal.butter(N=N, Wn=lowPass, btype='lowpass')
    filtSignal = signal.filtfilt(b, a, signalWav)  # signalWav为要过滤的信号
    return filtSignal

def HighFiltWav(signalWav, N, highPass):
    b, a = signal.butter(N=N, Wn=highPass, btype='highpass')
    filtSignal = signal.filtfilt(b, a, signalWav)  # signalWav为要过滤的信号
    return filtSignal

def ThisSTFT(filePath, sonicFreq, lipOffset, nfft, frameTime, aheadTime, saveDir):
    fs, signalWav = wav.read(filePath)

    fileName = filePath.split(os.sep)[-1]
    fileName = os.path.splitext(fileName)[0]
    savePath = os.path.join(saveDir, fileName)

    # 显示原始波形
    signalNum = len(signalWav)
    timeAxis = np.linspace(0, signalNum/fs, signalNum, endpoint=False)
    plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('amplitude')
    plt.plot(timeAxis, signalWav)
    plt.savefig(f'{savePath}_oriwav.jpg')
    plt.show()

    # 显示整个信号fft的图
    sonicPos = int(sonicFreq / fs * signalNum)
    offsetPoint = int(lipOffset / (fs / signalNum)) + 10
    fftSignal = np.fft.fft(signalWav)
    fftSignal = np.abs(fftSignal)
    freqAxis = np.linspace(0, fs, signalNum, endpoint=False)
    plt.title(filePath)
    plt.xlabel('freq(Hz)')
    plt.ylabel('amplitude')
    plt.plot(freqAxis[sonicPos - offsetPoint:sonicPos + offsetPoint],
             fftSignal[sonicPos - offsetPoint:sonicPos + offsetPoint])
    plt.savefig(f'{savePath}_fft.jpg')
    plt.show()


    lowPass = 2 * (sonicFreq - lipOffset) / fs
    highPass = 2 * (sonicFreq + lipOffset) / fs

    # 显示滤波后的波形
    filtSignal = FiltWav(signalWav=signalWav, N=4, lowPass=lowPass, highPass=highPass)
    plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('amplitude')
    plt.plot(timeAxis, filtSignal)
    plt.savefig(f'{savePath}_filtwav.jpg')
    plt.show()


    # filtSignal = filtSignal / max(abs(filtSignal))
    lmsSignal = filtSignal
    nperseg = int(fs * frameTime)
    noverlap = int(fs * (frameTime-aheadTime))
    f, t, Zxx = signal.stft(lmsSignal, fs, nperseg=nperseg, noverlap=noverlap, nfft=nfft)
    Zxx = Zxx / nfft * 2
    Zxx = CutSTFTBiside(Zxx)
    Zxx = np.abs(Zxx)
    startOff = 30
    t = t[startOff:Zxx.shape[1]+startOff]
    realOffset = lipOffset + 10
    down, up = LipGraphPara(sonicFreq, realOffset, fs, nfft)
    f = f[down:up]
    Zxx = Zxx[down:up,:]

    return t, f, Zxx

def SinTimeFreqShow(filePath, sonicFreq, lipOffset, nfft, frameTime, aheadTime, saveDir):
    t, f, Zxx = ThisSTFT(filePath, sonicFreq, lipOffset, nfft, frameTime, aheadTime, saveDir)
    # 显示短时傅里叶变换后的图
    # Zxx = np.log10(Zxx)

    fileName = filePath.split(os.sep)[-1]
    fileName = os.path.splitext(fileName)[0]
    savePath = os.path.join(saveDir, fileName)


    plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('freqency(hz)')
    plt.pcolormesh(t, f, Zxx)
    plt.savefig(f'{savePath}_stft.jpg')
    plt.show()
    # Zxx = np.log10(Zxx)
    # ShowSinFFT(Zxx, f)
    # 显示信号梯度图
    subZxx = SubAdjacent(Zxx)
    print(subZxx.shape)
    subZxx = np.abs(subZxx)
    subZxx = GaussSmooth2D(subZxx, (5, 5), 1)
    plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('freqency(hz)')
    plt.pcolormesh(t, f, subZxx)
    plt.savefig(f'{savePath}_grad.jpg')
    plt.show()

    # ShowSinFFT(subZxx, f, t)
    # 显示频率变化图
    t = t[:subZxx.shape[1]]
    dopplerShift = SimpleDopplerShift(subZxx, f)
    dopplerShift = GaussSmooth(dopplerShift, 5, 2)
    plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('frequency shift(hz)')
    plt.plot(t, dopplerShift)
    plt.savefig(f'{savePath}_freqchange.jpg')
    plt.show()
    # 显示频率变化归一化的图

def TimeFreqShow(corpus, sonicFreq, lipOffset, nfft, frameTime, aheadTime, saveDir):
    pattern = os.path.join(corpus, '*')
    for filePath in glob.glob(pattern):
        print(filePath)
        SinTimeFreqShow(filePath, sonicFreq, lipOffset, nfft, frameTime, aheadTime, saveDir)


def SinModem(filePath, sonicFreq, saveDir):
    fileName = filePath.split(os.sep)[-1]
    fileName = os.path.splitext(fileName)[0]
    savePath = os.path.join(saveDir, fileName)

    fs, signalWav = wav.read(filePath)
    N = 3
    highTh = 14000
    highPass = 2*highTh/fs
    signalWav = HighFiltWav(signalWav, N, highPass)
    signalLen = len(signalWav)
    time = np.linspace(0, signalLen / fs, signalLen, endpoint=False)

    coscw = np.cos(2 * np.pi * sonicFreq * time)
    sincw = -np.sin(2 * np.pi * sonicFreq * time)

    real = signalWav * coscw
    image = signalWav * sincw
    lowTh = 100
    lowPass = lowTh*2/fs

    real = LowFiltWav(real, N, lowPass)
    image = LowFiltWav(real, N, lowPass)

    startCut = int(0.3 * fs)
    endCut = int(0.12 * fs)

    real = real[startCut:len(real)-endCut]
    image = image[startCut:len(image)-endCut]
    time = time[startCut:len(time)-endCut]

    plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('amplitude')
    plt.plot(time, real)
    plt.savefig(f'{savePath}_modemreal.jpg')
    plt.show()

def Modem(corpus, sonicFreq, saveDir):
    pattern = os.path.join(corpus, '*')
    for filePath in glob.glob(pattern):
        print(filePath)
        SinModem(filePath, sonicFreq, saveDir)

if __name__ == '__main__':
    filePath = 'F:\\dataset\\lip_move\\easy_word_little'
    corpus = 'F:\\dataset\\lip_move\\easy_word_little'
    saveDir = 'F:\\dataset\\show_pic'
    sonicFreq = 20000
    lipOffset = 100
    nfft = 8192
    frameTime = 0.05
    aheadTime = 0.01
    TimeFreqShow(corpus, sonicFreq, lipOffset, nfft, frameTime, aheadTime, saveDir)
    # SinTimeFreqShow(filePath, sonicFreq, lipOffset, nfft, frameTime, aheadTime)
    # SinModem(filePath, sonicFreq, saveDir)
    # Modem(corpus, sonicFreq, saveDir)